package com.jason.job;

import com.jason.entity.DataInput;
import com.jason.mapper.DataMapper;
import com.jason.partitioner.DataPartitioner;
import com.jason.processor.DataPartitionProcessor;
import com.jason.reader.CommonPartitionMybatisItemReader;
import com.jason.writer.CommonMybatisItemWriter;
import org.apache.ibatis.session.SqlSessionFactory;
import org.springframework.batch.core.Job;
import org.springframework.batch.core.Step;
import org.springframework.batch.core.configuration.annotation.EnableBatchProcessing;
import org.springframework.batch.core.configuration.annotation.JobBuilderFactory;
import org.springframework.batch.core.configuration.annotation.StepBuilderFactory;
import org.springframework.batch.core.configuration.annotation.StepScope;
import org.springframework.batch.core.partition.PartitionHandler;
import org.springframework.batch.core.partition.support.TaskExecutorPartitionHandler;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.SimpleAsyncTaskExecutor;

/************************************************************************************************
 *  @Author:zhouhx
 *  @Date: 10:05 2020/11/18                                                                                                                                                 
 *  @Description: //分区处理job
 *  @Version:1.0
 ************************************************************************************************
 **/
@Configuration
@EnableBatchProcessing
public class DataPartitionerJob {

    @Autowired
    private JobBuilderFactory jobBuilderFactory;

    @Autowired
    private StepBuilderFactory stepBuilderFactory;

    @Autowired
    private DataPartitioner dataPartitioner;

    @Autowired
    private SqlSessionFactory sqlSessionFactory;

    @Autowired
    private DataPartitionProcessor dataPartitionProcessor;

    @Bean
    public Job dataPartitionJob() {
        return jobBuilderFactory.get("dataPartitionJob")
                .start(dataMasterStep())
                .build();
    }

    @Bean
    public Step dataMasterStep() {
        return stepBuilderFactory.get("dataMasterStep").partitioner(dataSlaveStep().getName(), dataPartitioner)
                .partitionHandler(dataPartitionHandler()).build();
    }

    @Bean
    public PartitionHandler dataPartitionHandler() {
        TaskExecutorPartitionHandler handler = new TaskExecutorPartitionHandler();
        handler.setGridSize(400);
        handler.setTaskExecutor(dataPartitionHandlerTaskExecutor());
        handler.setStep(dataSlaveStep());
        try {
            handler.afterPropertiesSet();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return handler;
    }

    @Bean
    public SimpleAsyncTaskExecutor dataPartitionHandlerTaskExecutor() {
        return new SimpleAsyncTaskExecutor();
    }

    @Bean
    public Step dataSlaveStep() {
        return stepBuilderFactory.get("dataSlaveStep")
                .<DataInput, DataInput>chunk(500)
                .reader(commonPartitionMybatisItemReader(null,null))
                .processor(dataPartitionProcessor)
                .writer(commonPartitionMybatisItemWriter())
                .build();
    }

    @Bean
    @StepScope
    public CommonPartitionMybatisItemReader commonPartitionMybatisItemReader(@Value("#{stepExecutionContext[fromId]}") final String fromId,
                                                                             @Value("#{stepExecutionContext[toId]}") final String toId) {
        return new CommonPartitionMybatisItemReader(sqlSessionFactory, DataInput.class.getSimpleName(),fromId,toId);
    }

    @Bean
    @StepScope
    public CommonMybatisItemWriter commonPartitionMybatisItemWriter() {
        return new CommonMybatisItemWriter(sqlSessionFactory, DataMapper.class.getSimpleName());
//        return new CommonItemWriter();
    }
}
